
import array
import sys
import Image
import time
import os

def load(fname):
    im = Image.open(fname)
    return im

class FatImg(object):
    def __init__(self, im, border_value, out_dir):
        self.im = im
        self.border_value = border_value
        self.out_dir = out_dir
        self._cached = {}

    def get(self, key):
        try:
            return self._cached[key]
        except KeyError:
            cmd = key.rsplit('_',1)[0]
            black, white = getattr(self, cmd)()
            self._cached[cmd+'_black'] = black
            self._cached[cmd+'_white'] = white
            assert key==(cmd+'_black') or key==(cmd+'_white') 
            return self._cached[key]
            
    def left(self):
        border_value = self.border_value
        im = self.im
        width, height = im.size
        black = array.array('h',( 0 for x in xrange(width*height))) #signed short 
        white = array.array('h',( 0 for x in xrange(width*height))) #signed short 
        for y in xrange(height):
            a = border_value
            b = border_value
            index = y*width
            for x in xrange(width):
                if im.getpixel((x,y))==0: #black:
                    a += 1
                    b = 0
                else:
                    b += 1
                    a = 0
                black[index] = a #left_black(x,y) = a
                white[index] = b #left_white(x,y) = b
                index += 1
                
        return black, white

    def right(self):
        border_value = self.border_value
        im = self.im
        width, height = im.size
        black = array.array('h',[0]*(width*height)) #signed short 
        white = array.array('h',[0]*(width*height)) #signed short 
        for y in xrange(height):
            a = border_value
            b = border_value
            for x in xrange(width-1,-1,-1):
                if im.getpixel((x,y))==0: #black
                    a += 1
                    b = 0
                else:
                    b += 1
                    a = 0
                black[y*width + x] = a 
                white[y*width + x] = b 
                
        return black, white

    def top(self):
        border_value = self.border_value
        im = self.im
        width, height = im.size
        black = array.array('h',( 0 for x in xrange(width*height))) #signed short 
        white = array.array('h',( 0 for x in xrange(width*height))) #signed short 
        for x in xrange(width):
            a = border_value
            b = border_value
            for y in xrange(height):
                if im.getpixel((x,y))==0: #black:
                    a += 1
                    b = 0
                else:
                    b += 1
                    a = 0
                black[y*width+x] = a 
                white[y*width+x] = b 
                
        return black, white

    def bottom(self):
        border_value = self.border_value
        im = self.im
        width, height = im.size
        black = array.array('h',( 0 for x in xrange(width*height))) #signed short 
        white = array.array('h',( 0 for x in xrange(width*height))) #signed short 
        for x in xrange(width):
            a = border_value
            b = border_value
            for y in xrange(height-1,-1,-1):
                if im.getpixel((x,y))==0: #black:
                    a += 1
                    b = 0
                else:
                    b += 1
                    a = 0
                black[y*width+x] = a 
                white[y*width+x] = b 
                
        return black, white

    def xaxis(self):
        forward = self.get('left_black')
        backward = self.get('right_black')
        black = array_min(forward, backward)
        forward = self.get('left_white')
        backward = self.get('right_white')
        white = array_min(forward, backward)
        return black, white

    def yaxis(self):
        forward = self.get('top_black')
        backward = self.get('bottom_black')
        black = array_min(forward, backward)
        forward = self.get('top_white')
        backward = self.get('bottom_white')
        white = array_min(forward, backward)
        return black, white

    def innerness_4way(self):
        xaxis = self.get('xaxis_black')
        yaxis = self.get('yaxis_black')
        black = array_min(xaxis, yaxis)
        xaxis = self.get('xaxis_white')
        yaxis = self.get('yaxis_white')
        white = array_min(xaxis, yaxis)
        return black, white
        

    def NE(self):
        border_value = self.border_value
        im = self.im
        width, height = im.size
        black = array.array('h',( 0 for x in xrange(width*height))) #signed short 
        white = array.array('h',( 0 for x in xrange(width*height))) #signed short 
        x0 = 0
        for y0 in xrange(height):
            a = border_value
            b = border_value
            x = x0
            y = y0
            while (y >= 0) and (x < width):
                if im.getpixel((x,y))==0: #black:
                    a += 1
                    b = 0
                else:
                    b += 1
                    a = 0
                black[y*width+x] = a 
                white[y*width+x] = b
                x += 1
                y += -1

        y0 = height-1            
        for x0 in xrange(1,width):
            a = border_value
            b = border_value
            x = x0
            y = y0
            while (y >= 0) and (x < width):
                if im.getpixel((x,y))==0: #black:
                    a += 1
                    b = 0
                else:
                    b += 1
                    a = 0
                black[y*width+x] = a 
                white[y*width+x] = b
                x += 1
                y += -1
        return black, white

    def SE(self):
        border_value = self.border_value
        im = self.im
        width, height = im.size
        black = array.array('h',( 0 for x in xrange(width*height))) #signed short 
        white = array.array('h',( 0 for x in xrange(width*height))) #signed short 
        y0 = 0
        for x0 in xrange(width):
            a = border_value
            b = border_value
            x = x0
            y = y0
            while (y < height) and (x < width):
                if im.getpixel((x,y))==0: #black:
                    a += 1
                    b = 0
                else:
                    b += 1
                    a = 0
                black[y*width+x] = a 
                white[y*width+x] = b
                x += 1
                y += 1

        x0 = 0            
        for y0 in xrange(1,height):
            a = border_value
            b = border_value
            x = x0
            y = y0
            while (y < height) and (x < width):
                if im.getpixel((x,y))==0: #black:
                    a += 1
                    b = 0
                else:
                    b += 1
                    a = 0
                black[y*width+x] = a 
                white[y*width+x] = b
                x += 1
                y += 1
        return black, white

    def NW(self):
        border_value = self.border_value
        im = self.im
        width, height = im.size
        black = array.array('h',( 0 for x in xrange(width*height))) #signed short 
        white = array.array('h',( 0 for x in xrange(width*height))) #signed short 
        x0 = width - 1
        for y0 in xrange(height):
            a = border_value
            b = border_value
            x = x0
            y = y0
            while (y >= 0) and (x >= 0):
                if im.getpixel((x,y))==0: #black:
                    a += 1
                    b = 0
                else:
                    b += 1
                    a = 0
                black[y*width+x] = a 
                white[y*width+x] = b
                x += -1
                y += -1

        y0 = height - 1            
        for x0 in xrange(width-1):
            a = border_value
            b = border_value
            x = x0
            y = y0
            while (y >= 0) and (x >= 0):
                if im.getpixel((x,y))==0: #black:
                    a += 1
                    b = 0
                else:
                    b += 1
                    a = 0
                black[y*width+x] = a 
                white[y*width+x] = b
                x += -1
                y += -1
        return black, white

    def SW(self):
        border_value = self.border_value
        im = self.im
        width, height = im.size
        black = array.array('h',( 0 for x in xrange(width*height))) #signed short 
        white = array.array('h',( 0 for x in xrange(width*height))) #signed short 
        y0 = 0            
        for x0 in xrange(width):
            a = border_value
            b = border_value
            x = x0
            y = y0
            while (y <height) and (x >= 0):
                if im.getpixel((x,y))==0: #black:
                    a += 1
                    b = 0
                else:
                    b += 1
                    a = 0
                black[y*width+x] = a 
                white[y*width+x] = b
                x += -1
                y += 1

        x0 = width-1            
        for y0 in xrange(height-1):
            a = border_value
            b = border_value
            x = x0
            y = y0
            while (y <height) and (x >= 0):
                if im.getpixel((x,y))==0: #black:
                    a += 1
                    b = 0
                else:
                    b += 1
                    a = 0
                black[y*width+x] = a 
                white[y*width+x] = b
                x += -1
                y += 1
        return black, white

    def NEaxis(self):
        forward = self.get('NE_black')
        backward = self.get('SW_black')
        black = array_min(forward, backward)
        forward = self.get('NE_white')
        backward = self.get('SW_white')
        white = array_min(forward, backward)
        return black, white

    def NWaxis(self):
        forward = self.get('NW_black')
        backward = self.get('SE_black')
        black = array_min(forward, backward)
        forward = self.get('NW_white')
        backward = self.get('SE_white')
        white = array_min(forward, backward)
        return black, white

    def innerness_4way_alt(self):
        one = self.get('NWaxis_black')
        two = self.get('NEaxis_black')
        black = array_min(one, two)
        one = self.get('NWaxis_white')
        two = self.get('NEaxis_white')
        white = array_min(one, two)
        return black, white

    def innerness_8way(self):
        one = self.get('innerness_4way_black')
        two = self.get('innerness_4way_alt_black')
        black = array_min(one, two)
        one = self.get('innerness_4way_white')
        two = self.get('innerness_4way_alt_white')
        white = array_min(one, two)
        return black, white


    def array_to_im_raw(self, a):
        width, height = self.im.size
        mina = min(a)
        maxa = max(a)
        d = maxa - mina
        if d==0:
            d = 1
        f = 255.0/float(d)
            
        im = Image.new('P', (width, height))
        for y in xrange(height):
            for x in xrange(width):
                im.putpixel((x,y),int((a[y*width+x]-mina)*f))
        return im

    def array_to_im_linear_fit(self, a):
        width, height = self.im.size
        im = Image.new('P', (width, height))
        for y in xrange(height):
            for x in xrange(width):
                im.putpixel((x,y),a[y*width+x])
        return im

    def array_to_im_log_fit(self, a):
        from math import log
        width, height = self.im.size
        mina = min(a)
        maxa = max(a)
        f = 255.1 / log(maxa - mina + 1.0)
            
        im = Image.new('P', (width, height))
        for y in xrange(height):
            for x in xrange(width):
                im.putpixel((x,y),int(f*log(a[y*width+x] - mina + 1.0) + 0.5))
        return im

    def im_to_a(self, im):
        #print 'im.mode:', im.mode
        assert isinstance(im.getpixel((0,0)), int)
        #assert im.mode=='P'
        width, height = im.size
        a = array.array('h',( 0 for x in xrange(width*height))) #signed short 
        #print im.getpixel((0,0))
        for y in xrange(height):
            i = y*width
            for x in xrange(width):
                a[i] = im.getpixel((x,y))
                i += 1
        return a
        
    def dump_4way(self, prefix):
        names = ['left_black', 'left_white', 'right_black', 'right_white',
                 'top_black', 'top_white', 'bottom_black', 'bottom_white',
                 'xaxis_black', 'xaxis_white', 'yaxis_black', 'yaxis_white',
                 'innerness_4way_black', 'innerness_4way_white' ]
        for name in names:
            a = self.get(name)
            im = self.array_to_im_log_fit(a)
            sname = prefix + name + '.png'
            fname = os.path.join(self.out_dir, sname)
            im.save(fname)

    def dump_diags(self, prefix):
        names = ['NE_black', 'NW_black', 'SW_black', 'SE_black',
                 'NE_white', 'NW_white', 'SW_white', 'SE_white',
                 'NWaxis_black','NEaxis_black', 'NWaxis_white','NEaxis_white',
                 'innerness_4way_alt_black', 'innerness_4way_alt_white' ]
        for name in names:
            a = self.get(name)
            im = self.array_to_im_log_fit(a)
            sname = prefix + name + '.png'
            fname = os.path.join(self.out_dir, sname)
            im.save(fname)

    def dump_all(self, prefix):
        self.dump_4way(prefix)
        self.dump_diags(prefix)
        names = ['innerness_8way_black', 'innerness_8way_white']
        for name in names:
            a = self.get(name)
            im = self.array_to_im_log_fit(a)
            sname = prefix + name + '.png'
            fname = os.path.join(self.out_dir, sname)
            im.save(fname)

    def load_all(self, prefix=None, src_getmtime=None):
        if src_getmtime is None:
            src_getmtime = 0 # WARN: will rebuilt
        # parts are rebuilt if not present or outdated
        names0 = ['left_black', 'left_white', 'right_black', 'right_white',
                 'top_black', 'top_white', 'bottom_black', 'bottom_white',
                 'xaxis_black', 'xaxis_white', 'yaxis_black', 'yaxis_white',
                 'innerness_4way_black', 'innerness_4way_white' ]
        names1 = ['NE_black', 'NW_black', 'SW_black', 'SE_black',
                 'NE_white', 'NW_white', 'SW_white', 'SE_white',
                 'NWaxis_black','NEaxis_black', 'NWaxis_white','NEaxis_white',
                 'innerness_4way_alt_black', 'innerness_4way_alt_white' ]
        names2 = ['innerness_8way_black', 'innerness_8way_white']
        names = names0 + names1 + names2
        got_them = True
        for name in names:
            sname = prefix + name + '.png'
            fname = os.path.join(self.out_dir, sname)
            if os.path.exists(fname) and os.path.getmtime(fname)>=src_getmtime:
                im = load(fname)
                a = self.im_to_a(im)
                self._cached[name] = a
            else:
                a = self.get(name)
                im = self.array_to_im_log_fit(a)
                im.save(fname)

    def paint1(self, color_for_black, color_for_white):
        cb = color_for_black
        cw = color_for_white
        im = self.im
        width, height = im.size
        im2 = Image.new('RGB', (width, height))
        black_variation = 0.3
        white_variation = 0.3
        # black lower intensity
        cbl0 = cb[0]*(1.0-black_variation)
        cbl1 = cb[1]*(1.0-black_variation)
        cbl2 = cb[2]*(1.0-black_variation)
        mb = black_variation/255.0
        # white lower intensity
        cwl0 = cw[0]*(1.0-white_variation)
        cwl1 = cw[1]*(1.0-white_variation)
        cwl2 = cw[2]*(1.0-white_variation)
        mw = white_variation/255.0
        black = self.get('innerness_8way_black')
        white = self.get('innerness_8way_white')
        
        for y in xrange(height):
            i = y*width
            for x in xrange(width):
                if black[i]:
                    f = (1.0-black_variation) + black[i]*mb
                    c = (int(cbl0*f), int(cbl1*f), int(cbl2*f))
                else:
                    f = (1.0-white_variation) + white[i]*mw
                    c = (int(cwl0*f), int(cwl1*f), int(cwl2*f))
                im2.putpixel((x,y), c)
                i += 1
        return im2

## PIL im modes (string):
## 1 (1-bit pixels, black and white, stored as 8-bit pixels)
## L (8-bit pixels, black and white)
## P (8-bit pixels, mapped to any other mode using a colour palette)
## RGB (3x8-bit pixels, true colour)
## RGBA (4x8-bit pixels, true colour with transparency mask)
## CMYK (4x8-bit pixels, colour separation)
## YCbCr (3x8-bit pixels, colour video format)
## I (32-bit integer pixels)
## F (32-bit floating point pixels)
## PIL also supports a few special modes, including
##    RGBX (true colour with padding) and
##    RGBa (true colour with premultiplied alpha).
    
def array_min(*arrays):
    """array_min(a1,...,ak) -> array / array[i]==min(a1[i],...,ak[i]) for all i

    a1,...,ak of class array.array must have same len and same element type"""
    typecode = arrays[0].typecode
    n = len(arrays[0])
    for a in arrays:
        assert typecode==a.typecode
        assert n==len(a)
    mina = array.array(typecode, (min( (a[i] for a in arrays) ) for i in xrange(n)))
    return mina

def array_mean(*arrays):
    """array_mean(a1,...,ak) -> array / array[i]==(a1[i]+...+ak[i])/k for all i

    a1,...,ak of class array.array must have same len and same element type"""
    typecode = arrays[0].typecode
    n = len(arrays[0])
    for a in arrays:
        assert typecode==a.typecode
        assert n==len(a)
    k = len(arrays)
    meana = array.array(typecode, (sum( (a[i] for a in arrays) )/k for i in xrange(n)))
    return meana

def get_reference_rect(im, ref_color):
    """get_reference_rect(im, ref_color) -> (left, top, right, bottom)

    Detects a reference rectangle, with all pixel in color ref_color.
    No pixel outside the reference rectangle can have that color.
    In other worlds, the maximal rectangle (left, top, right, bottom) such that
        rect pix(x,y)==ref_color : left <= x <= right and top <= y <= bottom
    """
    width, height = im.size
    y = int(height/2)
    x = 0
    while im.getpixel((x,y))!=ref_color:
        x += 1
    left = x
    while y>=0 and (im.getpixel((x,y))==ref_color):
        y -= 1
    top = y+1
    y = int(height/2)
    while y<height and (im.getpixel((x,y))==ref_color):
        y += 1
    bottom = y -1
    y = bottom
    assert im.getpixel((x,y))==ref_color
    while x<width and (im.getpixel((x,y))==ref_color):
        x += 1
    right = x - 1
    return left, top, right, bottom
    
def crop(im, rect):
    """crop(im, rect) -> im2 other image with pixels copied from the rect in im

    rect is left, top, right, bottom in PIL coordinates
    borders are included, ie
    pix(x,y) : left <= x <= right and top <= y <= bottom
    are all copied
    """
    #! PIL crop dont includes last row and column (at least PIL 1.1.6)
    left, top, right, bottom = rect
    right += 1
    bottom += 1
    region = im.crop((left, top, right, bottom))
    im2 = Image.new(im.mode, region.size)
    w,h = region.size
    im2.paste(region, (0,0,w,h))
    return im2

def conv1_bw( im, white_alias_color):
    # pixels with color white_alias_color will be white, all others
    # will be black
    width, height = im.size
    im2 = Image.new('1', im.size)
    for y in xrange(height):
        for x in xrange(width):
            im2.putpixel((x,y), im.getpixel((x,y))==white_alias_color)
    return im2

def area_covered(im, color):
    area = 0
    width, height = im.size
    for y in xrange(height):
        for x in xrange(width):
            if im.getpixel((x,y))==color:
                area += 1
    return area


